import java.util.*;
import java.util.concurrent.atomic.AtomicInteger;

public class AvlTress<T extends Comparable<T>> {
    Node<T> treeRoot;
    class Node<T extends Comparable<T>>{
        T val;
        Node leftChild;
        Node rightChild;
        int height;
        Node(T val, Node leftChild, Node rightChild,int height){
            this.val=val;
            this.leftChild=leftChild;
            this.rightChild=rightChild;
            this.height=1;
        }
        Node(T val) {
            this(val, null, null,0);
        }
    }

    /**
     * RR失去平衡，需要左旋操作
     *
     *         1
     *            \
     *               2
     *                  \
     *                     3
     *
     *    leftRotate(root) ====>
     *           2
     *       /      \
     *     1           3
     * @param root
     * @return
     */
    public Node<T> leftRotate(Node<T> root){
        Node<T> currentRoot = root;
        if(currentRoot==null)
            throw new RuntimeException("空树不能进行leftRotate操作！");
        Node newRoot = currentRoot.rightChild;
        currentRoot.rightChild = currentRoot.rightChild.leftChild;
        newRoot.leftChild = currentRoot;
        // 更新树的高度
        currentRoot.height = getNodeHeight(currentRoot);
        newRoot.height = getNodeHeight(newRoot);
        return newRoot;
    }

    public int getNodeHeight(Node<T> root){
        Node<T> currentRoot = root;
        if(currentRoot==null)
            return 0;
        if(currentRoot.leftChild!=null && currentRoot.rightChild!=null)
            return 1 + Math.max(currentRoot.leftChild.height,currentRoot.rightChild.height);
        else if (currentRoot.leftChild==null && currentRoot.rightChild==null)
            return 1;
        else if (currentRoot.leftChild!=null)
            return currentRoot.leftChild.height + 1;
        else
            return currentRoot.rightChild.height + 1;
    }
    /**
     * LL失去平衡，需要右旋操作
     *
     *         1
     *       /
     *     2
     *    /
     *  3
     *
     *    rightRotate(root) ====>
     *           2
     *       /      \
     *     3           1
     * @param root
     * @return
     */
    public Node<T> rightRotate(Node<T> root){
        Node<T> currentRoot = root;
        if(root==null)
            throw new RuntimeException("空树不能进行右旋转操作！");
        Node newRoot = currentRoot.leftChild;
        root.leftChild = currentRoot.leftChild==null?null:currentRoot.leftChild.rightChild;
        newRoot.rightChild = currentRoot;

        // 更新树的高度
        currentRoot.height = getNodeHeight(currentRoot);
        newRoot.height = getNodeHeight(newRoot);
        return newRoot;
    }

    /**
     * LR失去平衡，先左旋后右旋
     */
    public Node<T> leftRightRotate(Node<T> root){
        Node<T> currentRoot = root;
        if(currentRoot==null)
            throw new RuntimeException("空树不能进行leftRightRotate操作！");
        currentRoot.leftChild = leftRotate(currentRoot.leftChild);
        return rightRotate(currentRoot);
    }

    public Node<T> insert(T val){
        if(treeRoot==null){
            treeRoot = new Node<>(val);
            return treeRoot;
        }
        // 存储路径上的祖先节点
        Node<T> nodeResult = null;
        Stack<Node> stack = new Stack<>();
        Node<T> currentRoot = treeRoot;
        while (currentRoot!=null){
            // 记录父节点
            stack.push(currentRoot);
            if(currentRoot.val.compareTo(val)==0)
                return currentRoot;
            else if(currentRoot.val.compareTo(val)<0)
                currentRoot = currentRoot.rightChild;
            else
                currentRoot = currentRoot.leftChild;
        }
        // 插入新节点
        Node<T> top= stack.peek();
        nodeResult = new Node<T>(val);
        stack.push(nodeResult);
        if(top.val.compareTo(val)>0)
            top.leftChild = nodeResult;
        else
            top.rightChild = nodeResult;
        // 寻找需要调整的节点，并更新高度,滚动数组
        Node<T> p1=null, p2 = null,p3=null;
        while (!stack.isEmpty()){
            Node<T> currentNode = stack.pop();
            // 检查平衡因子
            if(Math.abs(getNodeHeight(currentNode.leftChild) - getNodeHeight(currentNode.rightChild)) <= 1){
                currentNode.height = 1 + Math.max(getNodeHeight(currentNode.leftChild),getNodeHeight(currentNode.rightChild));
                if(p2==null){
                    p2 = currentNode;
                }else{
                    // 滚动
                    p3 = p2;
                    p2 = currentNode;
                }
            }else{
                // 失去平衡，需要调整
                p1 = currentNode;
                Node<T> newTreeRoot=null;
                // 判断不平衡的类型
                if(p1.val.compareTo(p2.val)<0){
                    // R
                    if(p2.val.compareTo(p3.val)<0){
                        // RR
                        newTreeRoot =leftRotate(p1);
                    }else{
                        // RL
                        newTreeRoot = rightLeftRotate(p1);
                    }
                }else{
                    // L
                    if(p2.val.compareTo(p3.val)<0){
                        // LR
                        newTreeRoot = leftRightRotate(p1);
                    }else{
                        // LL
                        newTreeRoot =  rightRotate(p1);
                    }
                }
                // 是根节点
                if (stack.isEmpty()) {
                    treeRoot = newTreeRoot;
                }else{
                    Node<T> parent = stack.peek();
                    // 放到父节点中
                    if(parent.val.compareTo(newTreeRoot.val)<0){
                        parent.rightChild = newTreeRoot;
                    }else
                        parent.leftChild = newTreeRoot;
                }
            }
        }

        return nodeResult;
    }

    /**
     * RL失去平衡，先右旋后左旋
     */
    public Node<T> rightLeftRotate(Node<T> root){
        Node<T> currentRoot = root;
        currentRoot.rightChild = rightRotate(currentRoot.rightChild);
        return leftRotate(currentRoot);
    }

    // 搜索;             searchDeep：使用引用的方式返回本次搜索经历的搜索深度
    public Node<T> search(T val, AtomicInteger searchDeep){
        Node<T> currentRoot = treeRoot;
        while (currentRoot!=null){
            searchDeep.incrementAndGet();
            if(currentRoot.val.compareTo(val)==0)
                return currentRoot;
            else if(currentRoot.val.compareTo(val)<0)
                currentRoot = currentRoot.rightChild;
            else
                currentRoot = currentRoot.leftChild;
        }
        return null;
    }

    // 按行打印二叉树
    public List<List<Node<T>>> orderLevel(){
        ArrayList<List<Node<T>>> result = new ArrayList<>();
        Deque<Node<T>> queue = new ArrayDeque<>();
        Node<T> currentRoot = treeRoot;
        queue.addLast(currentRoot);
        while (!queue.isEmpty()){
            int levelSize = queue.size();
            ArrayList<Node<T>> levelNodes = new ArrayList<>(levelSize);
            for (int i = 0; i < levelSize; i++) {
                Node<T> topNode = queue.removeFirst();
                // visit
                levelNodes.add(topNode);
                if (topNode.leftChild!=null)
                    queue.addLast(topNode.leftChild);
                if(topNode.rightChild!=null)
                    queue.addLast(topNode.rightChild);
            }
            result.add(levelNodes);
        }
        return result;
    }
}
